# src/utils/db_utils.py
import mysql.connector
from mysql.connector import Error
from core.models import RawData, UserSelection, FieldSelection
from typing import List, Optional, Dict, Any
import logging
import re

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class DatabaseManager:
    def __init__(self):
        """初始化数据库连接"""
        try:
            self.connection = mysql.connector.connect(
                host='localhost',
                user='root',
                password='123456',
                database='well_data',
                charset='utf8mb4',
                autocommit=True,
                connection_timeout=10
            )
            if self.connection.is_connected():
                logger.info('Connected to MySQL database')
                self.create_tables()  # 自动创建数据表
        except Error as e:
            logger.error(f"Error while connecting to MySQL: {e}", exc_info=True)
            raise

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            logger.info("Database connection closed")

    def execute_query(self, query, params=None):
        """执行查询操作"""
        with self.connection.cursor(dictionary=True) as cursor:
            try:
                cursor.execute(query, params or ())
                result = cursor.fetchall()
                logger.debug(f"Query executed successfully: {query}")
                return result
            except Error as e:
                logger.error(f"Query execution failed: {query}, Error: {e}", exc_info=True)
                raise

    def execute_update(self, query, params=None):
        """执行更新操作"""
        with self.connection.cursor() as cursor:
            try:
                cursor.execute(query, params or ())
                logger.debug(f"Update executed successfully: {query}")
                return True
            except Error as e:
                logger.error(f"Update failed: {query}, Error: {e}", exc_info=True)
                return False

    def create_tables(self):
        """创建数据库表结构"""
        queries = [
            """
            CREATE TABLE IF NOT EXISTS raw_data (
                id INT AUTO_INCREMENT PRIMARY KEY,
                well_name VARCHAR(255) NOT NULL UNIQUE,
                base_level INT DEFAULT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """,

            """
            CREATE TABLE IF NOT EXISTS field_metadata (
                field_id INT AUTO_INCREMENT PRIMARY KEY,
                field_name VARCHAR(255) NOT NULL UNIQUE,
                table_name VARCHAR(255) NOT NULL UNIQUE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """,

            """
            CREATE TABLE IF NOT EXISTS user_selections (
                selection_id INT AUTO_INCREMENT PRIMARY KEY,
                raw_id INT NOT NULL,
                start_row INT NOT NULL,
                end_row INT NOT NULL,
                selected_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (raw_id) REFERENCES raw_data(id) ON DELETE CASCADE,
                CONSTRAINT valid_row_range CHECK (end_row >= start_row)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """,

            """
            CREATE TABLE IF NOT EXISTS selection_fields (
                selection_id INT NOT NULL,
                field_id INT NOT NULL,
                alias VARCHAR(255),
                is_visible BOOLEAN DEFAULT TRUE,
                PRIMARY KEY (selection_id, field_id),
                FOREIGN KEY (selection_id) REFERENCES user_selections(selection_id) ON DELETE CASCADE,
                FOREIGN KEY (field_id) REFERENCES field_metadata(field_id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """
        ]

        for query in queries:
            try:
                with self.connection.cursor() as cursor:
                    cursor.execute(query)
                logger.info(f"Table created/verified successfully: {query.split(' ')[5]}...")
            except Error as e:
                logger.error(f"Table creation failed: {e}", exc_info=True)
                raise

    # ================= 行号操作 =================
    def insert_raw_data(self, well_name: str, base_level: Optional[int] = None) -> Optional[int]:
        """插入原始数据记录（支持 base_level）"""
        query = """
            INSERT INTO raw_data (well_name, base_level) 
            VALUES (%s, %s)
        """
        with self.connection.cursor() as cursor:
            try:
                cursor.execute(query, (well_name, base_level))
                return cursor.lastrowid
            except Error as e:
                logger.error(f"行记录插入失败: {e}", exc_info=True)
                return None

    def get_raw_metadata(self, raw_id: int) -> Optional[RawData]:
        """获取原始数据元信息"""
        query = """
            SELECT id, well_name, base_level, created_at 
            FROM raw_data WHERE id = %s
        """
        result = self.execute_query(query, (raw_id,))
        return RawData(**result[0]) if result else None

    # ================= 字段操作 =================
    def create_field_table(self, field_name: str) -> Optional[str]:
        """
        创建字段表（显式存储行号）
        """
        table_name = self._validate_table_name(field_name)
        print(f"创建字段表: {table_name}")

        # 1. 检查是否已存在
        check_query = """
            SELECT table_name FROM field_metadata 
            WHERE field_name = %s OR table_name = %s
        """
        result = self.execute_query(check_query, (field_name, table_name))
        print(f"检查结果: {result}")

        # 2. 如果存在，返回对应表名
        if result:
            existing_table = result[0]['table_name']
            if existing_table == table_name:
                print(f"字段 {field_name} 已存在，表名匹配: {table_name}")
                return table_name
            else:
                logger.warning(f"字段名冲突: {field_name} | 现有表名: {existing_table}")
                return None  # 字段名与现有表名不匹配时返回 None

        print("创建新表")

        # 3. 检查连接状态
        if not self.connection or self.connection.is_closed():
            logger.error("数据库连接无效或已关闭")
            raise ConnectionError("无效数据库连接")

        print(f"当前连接: {self.connection}")

        # 4. 构建 SQL 查询
        create_query = f"""
            CREATE TABLE IF NOT EXISTS `{table_name}` (
                id INT AUTO_INCREMENT PRIMARY KEY,
                raw_id INT NOT NULL,
                `row_number` INT NOT NULL,
                col_value DECIMAL(10,2) NOT NULL,
                FOREIGN KEY (raw_id) REFERENCES raw_data(id) ON DELETE CASCADE,
                UNIQUE (raw_id, `row_number`)  -- 这里添加反引号
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """

        print(f"正在执行的 SQL 语句: {create_query}")  # 调试日志

        try:
            with self.connection.cursor() as cursor:
                # 5. 执行 SQL 查询
                res = cursor.execute(create_query)
                print(f"创建表 SQL 响应: {res}")

                # 6. 插入元数据
                metadata_query = """
                    INSERT INTO field_metadata (field_name, table_name)
                    VALUES (%s, %s)
                """
                res = cursor.execute(metadata_query, (field_name, table_name))
                print(f"插入元数据 SQL 响应: {res}")

                print(f"字段表创建成功: {table_name}...")
                return table_name

        except Error as e:
            logger.error(f"SQL 执行失败: {e}")
            logger.error(f"SQL 语句: {create_query}")

            # 7. 更详细的错误诊断
            if "1064" in str(e):  # MySQL 错误码 1064: 语法错误
                logger.error("可能原因：SQL 语法错误，请检查语句结构")
            elif "1050" in str(e):  # 表已存在
                logger.error("表已存在，但未被 metadata 记录，请手动清理数据库")
            elif "1046" in str(e):  # 没有选择数据库
                logger.error("未指定数据库，请确认连接参数中的 database 参数")
            elif "1146" in str(e):  # 表不存在（外键引用失败）
                logger.error("外键引用失败，请确认 raw_data 表已存在")
            elif "1072" in str(e):  # 索引失败
                logger.error("UNIQUE 约束失败，请检查字段类型/长度")
            elif "1044" in str(e):  # 权限问题
                logger.error("MySQL 用户权限不足，请授予 CREATE 权限")
            elif "1215" in str(e):  # 外键约束失败
                logger.error("外键约束失败，请检查 raw_data 表是否存在且为 InnoDB 引擎")
            else:
                logger.error(f"未知错误: {e}", exc_info=True)

            return None
    def _validate_table_name(self, name: str) -> str:
        """校验并转换合法表名"""
        # 替换非法字符
        valid_name = re.sub(r'[^a-zA-Z0-9_]', '_', name)
        # valid_name = name

        # 检查保留字冲突
        reserved_words = {'select', 'from', 'where', 'order', 'group', 'user', 'table', 'limit'}
        if valid_name.lower() in reserved_words:
            valid_name = f"{valid_name}_field"

        # 添加前缀并确保长度合法
        if not valid_name.startswith('tb_'):
            valid_name = f"tb_{valid_name}"

        # 限制表名长度（MySQL 最大表名长度为 64）
        if len(valid_name) > 64:
            valid_name = valid_name[:64]

        return valid_name

    def insert_data(self, table_name: str, raw_id: int, data: List[Dict[str, Any]]) -> bool:
        """
        批量插入字段值（显式传入行号）
        """
        if not table_name:
            logger.error("无效表名，插入失败")
            return False

        values = {}
        for i, row in enumerate(data):
            try:
                # 强制使用字段名获取值
                field_value = float(row.get(table_name.upper(), row.get(table_name.lower(), None)))
                values[i + 1] = field_value  # 行号 = 插入顺序
            except (KeyError, ValueError, TypeError):
                continue

        return self.batch_insert_field_values(table_name, raw_id, values)

    def batch_insert_field_values(self, table_name: str, raw_id: int, values: List[str]) -> bool:
        """
        批量插入字段值（使用 executemany 优化）
        """
        # 1. 验证表名合法性（防止 SQL 注入）
        if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', table_name):
            logger.error(f"非法表名: {table_name}")
            return False

        try:
            # 2. 数据预处理：过滤无效值并构建参数列表
            valid_params = []
            for row_num, value in enumerate(values):  # 行号从1开始
                try:
                    float_value = float(value)
                    valid_params.append((raw_id, row_num+1, float_value))
                    logger.debug(f"行 {row_num+1} 转换成功: {value} -> {float_value}")
                except (ValueError, TypeError) as e:
                    logger.warning(f"行 {row_num+1} 转换失败: {value}, 错误: {e}")

            if not valid_params:
                logger.warning("没有有效数据可插入")
                return False

            logger.info(f"准备批量插入 {len(valid_params)} 条记录到表 {table_name}")

            # 3. 构建 SQL 查询（单行插入模板）
            query = f"""
                INSERT INTO `{table_name}` 
                (raw_id, `row_number`, col_value) 
                VALUES (%s, %s, %s)
                ON DUPLICATE KEY UPDATE col_value = VALUES(col_value)
            """

            # 4. 执行批量插入
            with self.connection.cursor() as cursor:
                try:
                    cursor.executemany(query, valid_params)
                    row_count = cursor.rowcount
                    self.connection.commit()
                    logger.info(f"成功插入 {row_count} 条记录到表 {table_name}")
                    return True
                except mysql.connector.Error as e:
                    # 数据库错误处理
                    error_info = {
                        "error_code": e.errno,
                        "sql_state": e.sqlstate,
                        "message": str(e),
                        "query": query[:1000] + "..." if len(query) > 1000 else query,
                        "table": table_name,
                        "batch_size": len(valid_params)
                    }
                    logger.error(f"数据库错误: {error_info}", exc_info=True)
                    self.connection.rollback()
                    return False
                except Exception as e:
                    # 其他错误处理
                    logger.error(f"未知错误: {e}", exc_info=True)
                    self.connection.rollback()
                    return False

        except Exception as e:
            # 外层异常处理
            logger.error(f"批量插入过程异常: {e}", exc_info=True)
            if self.connection and self.connection.is_connected():
                self.connection.rollback()
            return False

    def get_field_data(self, table_name: str, raw_id: int, start_row: int, end_row: int) -> List[Dict]:
        """
        获取字段数据（使用 base_level 推导实际行号）
        """
        if not table_name:
            logger.error("无效表名，查询失败")
            return []

        # 获取 base_level
        raw_data = self.get_raw_metadata(raw_id)
        if not raw_data:
            return []

        base_level = raw_data.base_level or 0
        actual_start = start_row - base_level  # 确保行号 >= base_level + 1
        actual_end = end_row - base_level

        # 构建查询语句
        query = f"""
            SELECT row_number, col_value 
            FROM `{table_name}`
            WHERE raw_id = %s 
              AND `row_number` BETWEEN %s AND %s
            ORDER BY row_number
        """
        return self.execute_query(query, (raw_id, actual_start, actual_end))

    # ================= 用户选择操作 =================
    def insert_user_selection(self, raw_id: int, start_row: int, end_row: int) -> Optional[int]:
        """插入用户选择记录"""
        if end_row < start_row:
            logger.error("结束行号不能小于起始行号")
            return None

        query = """
            INSERT INTO user_selections 
            (raw_id, start_row, end_row) 
            VALUES (%s, %s, %s)
        """
        with self.connection.cursor() as cursor:
            try:
                cursor.execute(query, (raw_id, start_row, end_row))
                return cursor.lastrowid
            except Error as e:
                logger.error(f"插入用户选择失败: {e}", exc_info=True)
                return None

    def insert_selection_fields(self, selection_id: int, field_ids: List[int],
                               aliases: Optional[List[str]] = None,
                               is_visible_list: Optional[List[bool]] = None) -> bool:
        """插入选择的字段关联"""
        query = """
            INSERT INTO selection_fields 
            (selection_id, field_id, alias, is_visible) 
            VALUES (%s, %s, %s, %s)
        """

        values = []
        for i, field_id in enumerate(field_ids):
            alias = aliases[i] if aliases and i < len(aliases) else None
            is_visible = is_visible_list[i] if is_visible_list and i < len(is_visible_list) else True
            values.append((selection_id, field_id, alias, int(is_visible)))

        try:
            with self.connection.cursor() as cursor:
                cursor.executemany(query, values)
                return True
        except Error as e:
            logger.error(f"字段关联插入失败: {e}", exc_info=True)
            return False

    def get_user_selection(self, selection_id: int) -> Optional[UserSelection]:
        """获取完整用户选择对象"""
        query = """
            SELECT selection_id, raw_id, start_row, end_row, selected_at 
            FROM user_selections 
            WHERE selection_id = %s
        """
        result = self.execute_query(query, (selection_id,))
        if not result:
            return None

        selection_data = result[0]
        fields = self._load_selection_fields(selection_id)

        return UserSelection(
            selection_id=selection_data['selection_id'],
            raw_id=selection_data['raw_id'],
            start_row=selection_data['start_row'],
            end_row=selection_data['end_row'],
            selected_at=selection_data['selected_at'],
            fields=fields
        )

    def _load_selection_fields(self, selection_id: int) -> List[FieldSelection]:
        """加载字段关联信息"""
        query = """
            SELECT sf.field_id, sf.alias, sf.is_visible, fm.field_name 
            FROM selection_fields sf
            JOIN field_metadata fm ON sf.field_id = fm.field_id
            WHERE sf.selection_id = %s
        """
        results = self.execute_query(query, (selection_id,))

        return [
            FieldSelection(
                field_id=row['field_id'],
                field_name=row['field_name'],
                alias=row['alias'],
                is_visible=bool(row['is_visible'])
            ) for row in results
        ]

    # ================= 事务操作 =================
    def save_full_selection(self, raw_id: int, start_row: int, end_row: int,
                            fields: List[FieldSelection]) -> Optional[int]:
        """安全的保存方法"""
        if self.connection.in_transaction:
            logger.warning("事务已存在，跳过重复开启")

        try:
            self.connection.start_transaction()

            selection_id = self.insert_user_selection(raw_id, start_row, end_row)
            if not selection_id:
                raise Error("选择记录插入失败")

            field_ids = [f.field_id for f in fields]
            aliases = [f.alias for f in fields]
            is_visible_list = [f.is_visible for f in fields]

            if not self.insert_selection_fields(selection_id, field_ids, aliases, is_visible_list):
                raise Error("字段关联插入失败")

            self.connection.commit()
            return selection_id

        except Error as e:
            self.connection.rollback()
            logger.error(f"插入选择记录失败: {e}", exc_info=True)
            return None
